{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b739c0e2",
   "metadata": {},
   "source": [
    "# Image Detection Demo with PytorchWildlife"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1197e180",
   "metadata": {},
   "source": [
    "This tutorial guides you on how to use PyTorchWildlife to separate positive and negative animal detections. We will go through the process of setting up the environment, defining the detection model, as well as performing inference and saving the results in different ways.\n",
    "\n",
    "## Prerequisites\n",
    "Install PytorchWildlife running the following commands:\n",
    "```bash\n",
    "conda create -n pytorch_wildlife python=3.8 -y\n",
    "conda activate pytorch_wildlife\n",
    "pip install PytorchWildlife\n",
    "```\n",
    "Also, make sure you have a CUDA-capable GPU if you intend to run the model on a GPU. This notebook can also run on CPU.\n",
    "\n",
    "## Importing libraries\n",
    "First, we'll start by importing the necessary libraries and modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c44e7713",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from PytorchWildlife.models import detection as pw_detection\n",
    "from PytorchWildlife import utils as pw_utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6abd07b5",
   "metadata": {},
   "source": [
    "## Model Initialization\n",
    "We will initialize the MegaDetectorV5 model for image detection. This model is designed for detecting animals in images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb25db43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting the device to use for computations ('cuda' indicates GPU)\n",
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "if DEVICE == \"cuda\":\n",
    "    torch.cuda.set_device(0)\n",
    "    \n",
    "# Initializing the MegaDetectorV6 model for image detection\n",
    "# Valid versions are MDV6-yolov9-c, MDV6-yolov9-e, MDV6-yolov10-c, MDV6-yolov10-e or MDV6-rtdetr-c\n",
    "detection_model = pw_detection.MegaDetectorV6(device=DEVICE, pretrained=True, version=\"MDV6-yolov10-e\")\n",
    "\n",
    "# Uncomment the following line to use MegaDetectorV5 instead of MegaDetectorV6\n",
    "#detection_model = pw_detection.MegaDetectorV5(device=DEVICE, pretrained=True, version=\"a\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3cfe3fd",
   "metadata": {},
   "source": [
    "## Variable definition\n",
    "In order to process the batch detection, we will define an input directory where the images are stored, a confidence threshold and an output directory to copy the positive and negative images into distinctive folders. If you want to follow this tutorial with your own data, modify the following variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "735a3f71",
   "metadata": {},
   "outputs": [],
   "source": [
    "tgt_folder_path = os.path.join(\".\",\"demo_data\",\"imgs\")\n",
    "output_path = \"folder_separation\"\n",
    "threshold = 0.2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e23329c",
   "metadata": {},
   "source": [
    "## Batch Image Detection\n",
    "Next, we'll demonstrate how to process multiple images in batches. This is useful when you have a large number of images and want to process them efficiently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "561eff0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = detection_model.batch_image_detection(tgt_folder_path, batch_size=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb41830b",
   "metadata": {},
   "source": [
    "## Separate positive and negative detections\n",
    "PytorchWildlife allows to copy the files from your original folder to a new directory containing the \"Animal\" and \"No-animal\" subdirectories. A detection is considered positive if the prediction confidence is higher than the threshold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f63310ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(output_path, exist_ok=True)\n",
    "json_file = os.path.join(output_path, \"detection_results.json\")\n",
    "pw_utils.save_detection_json(results, json_file,\n",
    "                             categories=detection_model.CLASS_NAMES,\n",
    "                             exclude_category_ids=[], # Category IDs can be found in the definition of each model.\n",
    "                             exclude_file_path=tgt_folder_path)\n",
    "\n",
    "# Separate the positive and negative detections through file copying:\n",
    "pw_utils.detection_folder_separation(json_file, tgt_folder_path, output_path, threshold)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4ee1d7b",
   "metadata": {},
   "source": [
    "### Copyright (c) Microsoft Corporation. All rights reserved.\n",
    "### Licensed under the MIT License."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-wildlife",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
